"""

    File: flv.py
    Description: 
    
        This module contains infomation and functions needed to write or stream FLV files.

    Author: Kyle Vogt
    Date  : June 12th, 2007
    Copyright (c) 2007, Justin.tv, Inc.
    
"""

from twisted.internet import defer, threads
from twisted.python import threadable, log
threadable.init(1)
import time, struct, os, sys, random, traceback
from array import array
import traceback

def writeHeader():
    "Write the 13 byte FLV header"
    header = 'FLV\x01\x05\x00\x00\x00\t\x00\x00\x00\x00'
    return header
        
def thumb(tag, width=320, height=240, format='jpg'):
    "Use ffmpeg to generate a thumbnail preview from the given keyframe"
    def _():
        filler = 'f\xc2\x98/\xc6\x18c\x8c1\xc6\x18c\x8c1\xc6\xd8\xca\xd7\x97\x86\x9aVUU\x9aVVZ\xaa\x95oZ\xa5Z<\xf0\xff8\xfe\x18;fij\x9aYj\xda\xbai\xa6\x9a\xce\x9a\xba\x9aUjQ`\xf9\r\xe3\x96m'
        fake = {    'type' : 0x08, 
                    'data' : filler, 
                    'abstime' : 0, 
                    'timecode' : 0,
                    'object' : 3}
        data = writeHeader() + writeTag(tag, 0) + writeTag(fake, 0)
        temp = 'temp_%s_%s' % (int(time.time()), random.randint(0,10000))
        infile = temp + '.flv'
        outfile = temp + '.' + format
        h = open(infile, 'wb')
        h.write(data)
        h.close()
        command = "ffmpeg 2>>/dev/null -an -i %s -s %sx%s -vframes 1 -f image2 %s" % (infile, width, height, outfile)
        log.msg("command is: %s" % command)
        try:
            os.popen(command).read()
            f = open(outfile)
            data = f.read()
            f.close()
        except:
            log.msg('Error with FFMPEG:')
            traceback.print_exc()
        try:
            os.unlink(infile)
            os.unlink(outfile)
        except:
            log.msg('Thumbnail generator was unable to delete temp files %s and %s' % (infile, outfile))
        return data
    return threads.deferToThread(_)
    
def writeTag(tag, timecode):
    "Convert a tag into an FLV tag"
    if timecode < 0: timecode = 0
    data = chr(tag['type']) + struct.pack("!I", len(tag['data']))[1:]
    # Write the tag
    tstamp = struct.pack("!I", int(timecode))
    data += str(tstamp[1:])
    data += str(tstamp[0])
    data += chr(0) * 3 + tag['data']
    data += struct.pack("!l", int(len(tag['data']) + 11))
    return data
        
def writeMetaTag(metadata):
    "Wrap raw metadata in an FLV tag and return it"
    data = chr(0x12) + struct.pack("!I", len(metadata))[1:]
    data += chr(0x00) * 4
    data += chr(0) * 3 + metadata
    data += struct.pack("!l", int(len(metadata) + 11))
    return data
    
def readTag(data, offset = 0):
    "Read an FLV tag from a data buffer"
    try:
        if len(data[offset:]) < 8: return False
        tag_type = ord(data[offset])
        length = struct.unpack("!I", chr(0x00) + data[offset + 1: offset + 4])[0]
        time_stamp = struct.unpack("!I", data[offset + 7] + data[offset + 4: offset + 7])[0]
        if len(data[offset:]) >= 11 + length:
            tag_data = data[offset + 11 : offset + 11 + length]
            return {'data' : tag_data, 'type' : tag_type, 'time_stamp' : time_stamp}
        else: return False
    except:
        log.msg('Error reading FLV tag')
        return False
    
def readFile(file_name):
    "Parse an FLV file into a list of tags"
    log.msg('Reading file %s' % file_name)
    h = None
    try:
        h = open(file_name, 'rb')
        if h.read(13)[:3] != 'FLV': 
            log.msg('Bad FLV header!')
            return False
        data = h.read(65536) # use h.read(65536) instead for Mac's
    except:
        log.msg('FLV parser is unable to read file %s' % file_name)
        return False
    pointer = 0
    tags = []
    while True:
        tag = readTag(data, pointer)
        if tag:
            tag['byte_offset'] = 13 + pointer
            tag['key'] = isKeyframe(tag)
            tags.append(tag)
            pointer += 15 + len(tag['data'])
        else:
            newdata = h.read(65536)
            if newdata: data += newdata
            else: break
    try: h.close()
    except: pass
    return tags

def isKeyframe(tag):
    "Boolean fucntion to determine whether the FLV tag in this tag is a keyframe"
    return ((len(tag['data']) > 0) and (tag['type'] == 0x09) and getVideoFrame(tag) == 'keyframe')
        
def getVideoCodec(tag):
    "Determines whether the FLV tag in this tag uses the VP6 or Sorenson codec"
    info = ord(tag['data'][0])
    # Determine Codec type
    codec = (info & 0x0F)
    if codec == 2: return "sor"
    if codec == 4: return "vp6"
    else: return "err"
    
def getVideoFrame(tag):
    "Determines whether the FLV tag is a keyframe or not"
    frame = (ord(tag['data'][0]) & 0xF0) >> 4
    if frame == 1: return 'keyframe'
    elif frame == 2: return 'iframe'
    elif frame == 3: return 'unsupported'
    else: return 'err'
    
def getAudioCodec(tag):
    "Determines the audio codec of the given tag"
    format = (ord(tag['data'][0]) & 0xF0) >> 4
    if format == 0: return 'uncompressed'
    elif format == 1: return 'adpcm'
    elif format == 2: return 'mp3'
    elif format == 5: return 'nellymoser_8kz_mono'
    elif format == 6: return 'nellymoser'
    else: return 'err'
        